import PyNvCodec as nvc
import time
import numpy as np
import ctypes


def process_rtsp_stream(rtsp_url):
    # 设置RTSP流
    nvdec = nvc.PyNvDecoder(rtsp_url, 0, {"rtsp_transport": "tcp"})

    frame_count = 0
    total_decode_time = 0.0

    while True:
        frame = np.ndarray(shape=(0), dtype=np.uint8)
        raw_frame = nvc.PacketData()

        # 计算解码时间
        start_time = time.time()
        success = nvdec.DecodeSingleFrame(frame, raw_frame)
        decode_time = time.time() - start_time

        if not success:
            print("无法解码帧或流结束")
            break

        total_decode_time += decode_time
        frame_count += 1

        print(f"Frame {frame_count}: GPU Decode time = {decode_time:.6f} seconds")

        # 如果只处理前100帧，可以取消下面的注释
        # if frame_count >= 100:
        #     break

    print(f"Average decode time: {total_decode_time / frame_count:.6f} seconds")


# 替换为你的RTSP流地址
rtsp_url = "rtsp://127.0.0.1:8554/stream"
process_rtsp_stream(rtsp_url)
